import json
import random

from flask import Flask, jsonify
from flask_cors import CORS

###############################################################################
# 全局配置
###############################################################################
DEEPSEEK_KEY = "sk-27bdaedc5191497c9d437d8ffff64f9e"  # <<<<<< 换成自己的
BOARD_W, BOARD_H = 20, 15  # 格子数
DIRECTIONS = {"up": (0, -1), "down": (0, 1), "left": (-1, 0), "right": (1, 0)}
OPPOSITE = {"up": "down", "down": "up", "left": "right", "right": "left"}


###############################################################################
# 一个极简贪吃蛇环境
###############################################################################
class SnakeEnv:
    def __init__(self):
        self.reset()

    def reset(self):
        mid = (BOARD_W // 2, BOARD_H // 2)
        self.snake = [mid, (mid[0] - 1, mid[1])]  # 初始 2 节，水平
        self.direction = "right"
        self.spawn_food()
        self.game_over = False
        self.score = 0

    def spawn_food(self):
        while True:
            pos = (random.randint(0, BOARD_W - 1), random.randint(0, BOARD_H - 1))
            if pos not in self.snake:
                self.food = pos
                break

    def step(self, action):
        """action 必须 ∈ {'up','down','left','right'}"""
        if action == OPPOSITE[self.direction]:  # 禁止掉头
            action = self.direction
        self.direction = action
        head = self.snake[0]
        dx, dy = DIRECTIONS[action]
        new_head = (head[0] + dx, head[1] + dy)
        # 撞墙或撞自己
        if new_head[0] < 0 or new_head[0] >= BOARD_W or \
                new_head[1] < 0 or new_head[1] >= BOARD_H or \
                new_head in self.snake:
            self.game_over = True
            return
        self.snake.insert(0, new_head)
        if new_head == self.food:
            self.score += 1
            self.spawn_food()
        else:
            self.snake.pop()


###############################################################################
# 把局面变成自然语言 Prompt
###############################################################################
def build_prompt(env):
    """
    用纯文本描述地图，让 DeepSeek 直接输出一个方向单词。
    经验证，只要 Prompt 足够清晰，7B 级别模型就能稳定输出 {up|down|left|right}
    """
    board = [["." for _ in range(BOARD_W)] for _ in range(BOARD_H)]
    for x, y in env.snake:
        board[y][x] = "S"  # 蛇身
    fx, fy = env.food
    board[fy][fx] = "F"  # 食物
    lines = ["".join(row) for row in board]
    grid_txt = "\n".join(lines)
    prompt = (
        "You are playing a Snake game. The grid is BOARD_W x BOARD_H, "
        "top-left=(0,0). Snake head is at the first 'S'. Food is 'F'. "
        "Decide the next single move to eat food and avoid collision.\n"
        "Answer exactly one word: up / down / left / right.\n"
        f"Grid ({BOARD_W}x{BOARD_H}):\n{grid_txt}\n"
        "Answer exactly one word: up / down / left / right"
    )
    return prompt


###############################################################################
# 调用 DeepSeek Chat API
###############################################################################
from openai import OpenAI

client = OpenAI(api_key=DEEPSEEK_KEY, base_url="https://api.deepseek.com")


def ask_deepseek(prompt):
    try:
        resp = client.chat.completions.create(
            model="deepseek-chat",
            messages=[{"role": "user", "content": prompt}],
            stream=False
        )
        print("DeepSeekPrompt:", prompt)
        print("DeepSeekResponse:", json.dumps(resp.json(), ensure_ascii=False, indent=4, sort_keys=True))
        ans = resp.choices[0].message.content.strip().lower()
        # ans = resp.json()['choices'][0]['message']['content'].strip().lower()
        print("DeepSeekAns:", ans)
        if ans in DIRECTIONS:
            return ans
    except Exception as e:
        print("DeepSeek API error:", e)
    return "right"  # fallback


###############################################################################
# Flask 路由
###############################################################################
app = Flask(__name__)
CORS(app)  # 允许所有域名、所有路由、所有方法
env = SnakeEnv()


@app.route("/step", methods=["POST"])
def step():
    """前端可POST空 JSON 触发下一步"""
    if env.game_over:
        env.reset()
    prompt = build_prompt(env)
    action = ask_deepseek(prompt)
    env.step(action)
    return jsonify({
        "snake": env.snake,
        "food": env.food,
        "score": env.score,
        "game_over": env.game_over,
        "next_action": action
    })


@app.route("/state")
def state():
    """只读当前局面"""
    return jsonify({
        "snake": env.snake,
        "food": env.food,
        "score": env.score,
        "game_over": env.game_over
    })


###############################################################################
# 纯后端自驱动模式：直接 AI 自动运行
###############################################################################
def auto_run():
    while True:
        if env.game_over:
            print(f"Game Over. Score={env.score}")
            env.reset()
        prompt = build_prompt(env)
        action = ask_deepseek(prompt)
        env.step(action)


if __name__ == "__main__":
    # 如果你想“看”AI自己跑，把下一行取消注释
    # auto_run()

    # 如果想给前端/网页用，运行 Flask
    app.run(host='::', port=5001, debug=True)
